import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv,GINConv, global_add_pool

class GNNModel(nn.Module):
    def __init__(self, in_dim, hidden_dim=128, out_dim=2, num_layers=6, dropout=0.3):
        super().__init__()

        self.convs = nn.ModuleList()
        self.bns = nn.ModuleList()

        # 输入层
        self.convs.append(GINConv(nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )))
        self.bns.append(nn.BatchNorm1d(hidden_dim))

        # 中间层（包含残差连接）
        for _ in range(num_layers - 2):
            self.convs.append(GINConv(nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, hidden_dim)
            )))
            self.bns.append(nn.BatchNorm1d(hidden_dim))

        # 输出层
        self.fc = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(hidden_dim // 2, out_dim)
        )

        self.dropout = dropout
        self.residual = True  # 启用残差连接

    def forward(self, x, edge_index, batch=None):
        h = x
        for i, (conv, bn) in enumerate(zip(self.convs, self.bns)):
            h_new = conv(h, edge_index)
            h_new = bn(h_new)
            if self.residual and i > 0:  # 从第二层开始加残差
                h_new += h
            h = nn.functional.relu(h_new)
            h = nn.functional.dropout(h, p=self.dropout, training=self.training)

        # 图级池化
        if batch is not None:
            h = global_add_pool(h, batch)
        return self.fc(h)